import cv2
import numpy as np


def pose_calculate(K, X, x):
    distCoeffs = np.zeros(4)
    # 计算旋转向量和平移向量
    try:
        success, rvec, tvec = cv2.solvePnP(X, x, K, distCoeffs)
        if not success:
            Warning("Failed to solve PnP")
            return None, None
        # 将旋转向量转换为旋转矩阵
        R, _ = cv2.Rodrigues(rvec)
        return R, tvec
    except:
        return None, None


def ransac_pnp(K, points_3d, observations, th=8):
    best_num = 0
    best_inliers = np.array([False] * len(observations))
    sample = np.random.choice(len(observations), 4, replace=False)
    for _ in range(100):
        R, tvec = pose_calculate(K, points_3d[sample], observations[sample])
        if R is not None:
            err = np.linalg.norm(
                cv2.projectPoints(points_3d, R, tvec, K, np.zeros(4))[0].squeeze()
                - observations,
                axis=1,
            )
            inliers = err < th
            if inliers.sum() > best_num:
                best_num = inliers.sum()
                best_inliers = inliers
            if inliers.sum() > 0.8 * len(observations):
                break
            # 另一种方式，使用绝对点个数，实际上超过20个点后计算所得精度已经相当高了。
            # if inliers.sum() > 20:
            #     break
            if inliers.sum() > 4:
                sample = np.random.choice(np.where(inliers)[0], 4, replace=False)
            else:
                sample = np.random.choice(len(observations), 4, replace=False)
        else:
            sample = np.random.choice(len(observations), 4, replace=False)
    if best_inliers.sum() < 4:
        return None
    else:
        return best_inliers
